import torch
import torch.nn as nn
import torch.nn.functional as F

class MLPWithPooling(nn.Module):
    def __init__(self, input_size, output_size):
        super(MLPWithPooling, self).__init__()
        # 定义MLP的全连接层
        self.fc1 = nn.Linear(3 * input_size * input_size, 1000)  # 输入是展平的三通道64x64图像
        self.fc2 = nn.Linear(1000, 500)
        self.fc3 = nn.Linear(500, 3 * output_size * output_size)  # 输出是展平的三通道32x32图像
        self.pool = nn.MaxPool2d(kernel_size=1)  # 使用二维最大池化层
        self.out = output_size

    def forward(self, x):
        # 将输入张量展平，假设输入是(batch_size, 3, 64, 64)
        x = x.view(x.size(0), -1)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        # 将输出张量重新变形为(batch_size, 3, 32, 32)
        x = x.view(-1, 3, self.out, self.out)
        # 应用二维最大池化层
        x = self.pool(x)
        return x

class LSTMPool(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, out_size):
        super(LSTMPool, self).__init__()
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True)
        self.pool = nn.AdaptiveAvgPool2d((out_size , out_size))  # 设置自适应平均池化的目标尺寸为 (32, 32)

    def forward(self, x):
        batch_size, channels, height, width = x.size()
        x = x.view(batch_size, channels, height * width)  # 将图像转换为时间序列数据
        out, _ = self.lstm(x)
        out = out.permute(0, 2, 1)  # 将时间维度放回通道维度的后面，适应池化层的输入要求
        out = self.pool(out)  # 进行自适应平均池化，输出形状为 [batch_size, hidden_size, 32, 32]
        return out


class CNN_model(nn.Module):
    def __init__(self, Kernel_size1, Stride1, Kernel_size2, Stride2):
        super(CNN_model, self).__init__()

        self.conv1 = nn.Conv2d(in_channels=3, out_channels=3, kernel_size=Kernel_size1,  stride=Stride1).cuda()
        # if pool_type == 'max':
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2).cuda()
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        # elif pool_type == 'avg':
        #     self.pool1 = nn.AvgPool2d(kernel_size=2, stride=2).cuda()
        #     self.pool2 = nn.AvgPool2d(kernel_size=2, stride=2).cuda()
        # elif pool_type == 'lstm':
        #     self.pool1 = LSTMPool(64 * 64, 50, 1, 32)
        #     self.pool2 = LSTMPool(16 * 16, 50, 1, 8)
        # elif pool_type == 'mlp':
        #     self.pool1 = MLPWithPooling(64, 32)
        #     self.pool2 = MLPWithPooling(16, 8)

        self.conv2 = nn.Conv2d(in_channels=3, out_channels=3, kernel_size=Kernel_size2, stride=Stride2).cuda()
        self.dropout = nn.Dropout(p=0.7)
        self.fc1 = nn.Linear(192, 128).cuda()
        self.fc2 = nn.Linear(128, 77).cuda()
        # self.ex_fc = nn.Linear(77, 100).cuda()
        # self.ex_fc1 = nn.Linear(100, 100).cuda()
        # self.ex_fc2 = nn.Linear(100, 77).cuda()


    def forward(self, x):
        # x = x.unsqueeze(0)
        x = F.relu(self.conv1(x))
        x = self.pool1(x)
        x = F.relu(self.conv2(x))
        x = self.pool2(x)
        x = x.view(x.size(0), -1)
        x = F.log_softmax(x, dim=1)
        x = self.dropout(x)
        x = torch.sigmoid(self.fc1(x))
        x = torch.sigmoid(self.fc2(x))
        return x
        # outs = torch.sigmoid(self.ex_fc(x))
        # outs = self.dropout(outs)
        # outs = torch.sigmoid(self.ex_fc1(outs))
        # outs = self.dropout(outs)
        # outs = torch.sigmoid(self.ex_fc2(outs))
        # outs = self.dropout(outs)
        # return x, outs